# 通用的utils工具

from scipy.signal import resample
import numpy as np
from huggingface_hub import snapshot_download
import scipy
import os
import pandas as pd
import mat73
import pickle
from tqdm import tqdm
from multiprocessing import Pool
import json



# 下载模型到./pretrained_models/目录下，方便加载
def pretrain_model_download(repo_id, local_dir="./pretrained_models/"):
    snapshot_download(repo_id, local_dir=local_dir+repo_id, ignore_patterns=["*.msgpack", "*.h5", "*.onnx", "*.tflite", "*.ot", "*.safetensors"])


# 对eeg信号进行上采样。whisper模型的默认采样率为16000Hz
def up_sample_eeg(eegsignal, current_sample_rate, target_sample_rate=16000):
    eegsignal = np.array([eegsignal])  # 将您的音频数据列表转换为NumPy数组
    upsampling_factor = target_sample_rate / current_sample_rate # 计算上采样因子
    upsampled_eeg = resample(eegsignal, int(len(eegsignal) * upsampling_factor)) # 使用scipy的resample函数上采样
    return upsampled_eeg


# 把下载至rawdata的原始数据处理成可以直接使用data_loader加载的形式，dsname表示数据集的名称
# zuco     https://osf.io/q3zws/wiki/home/
# zuco2    https://osf.io/2urht/wiki/
def data_prepare(dsname):
    if dsname == "zuco":
        pass
    if dsname =="zuco2":
        data_prepare_zuco2()
    if dsname == "zuco3":
        pass

class Func4MultiProcessing: # 把多线程函数的逻辑部分打包进这个class，方便统一接口。
    def zuco2eeg(self, zuco2path, task_name, sbj_name, tn):
        # file_name = f"{zuco2path}/{task_name}/Matlab files/results{sbj_name}_{tn}.mat"
        file_name = f"./data/zuco2_{tn}.json"
        print(f"---==Reading[{file_name}]==---")
        # text_ = mat73.loadmat(file_name)['sentenceData']['content']
        text_ = json.load(open(file_name, "r", encoding="utf-8"))


        folder = f"{zuco2path}/{task_name}/Raw data/{sbj_name}/"
        filename_li = [f"{sbj_name}_{tn}{_}_EEG.mat" for _ in range(1, 8)]
        eeg_seg = []
        for file_name in tqdm(filename_li):
            print(f"---==Processing[{folder + file_name}]==---")
            data = scipy.io.loadmat(folder + file_name, squeeze_me=True, struct_as_record=False)
            raweeg = data["EEG"].data
            if "EEG" in data.keys():
                event = data["EEG"].event
            else:
                event = data["event"]


            cnt_10, cnt_11, cnt_12, cnt_13 = None, None, None, None
            cnt = 0
            for each in event:
                if each.value == "trigger":
                    if each.type.strip() == "10":
                        cnt_10 = each.latency
                    elif each.type.strip() == "11":
                        cnt_11 = each.latency
                        eeg_seg.append(raweeg[:, cnt_10:cnt_11+1])
                        cnt_10, cnt_11 = None, None
                        cnt += 1
                    elif each.type.strip() == "12":
                        cnt_12 = each.latency
                    elif each.type.strip() == "13":
                        cnt_13 = each.latency
                        eeg_seg.append(raweeg[:, cnt_12:cnt_13+1])
                        cnt_12, cnt_13 = None, None
                        cnt += 1
        assert len(eeg_seg) == len(text_)
        tmpdic = {"eeg_seg": eeg_seg, "text_": text_}
        pickle.dump(tmpdic, open(f"./tmp/zuco2_{task_name}_{sbj_name}.pkl", "wb"))
        print(f"---==Save [./tmp/zuco2_{task_name}_{sbj_name}.pkl]==---")

    @staticmethod
    def error_back(eb):
        print(f'error: {str(eb)}')


def mutipro_start(func, zuco2path, task_name, sbj_names_, tn):
    working_thread = []
    workernum = len(sbj_names_)
    with Pool(processes=workernum) as pool:
        for sbj_name in sbj_names_:
            w = pool.apply_async(
                func=func.zuco2eeg,
                args=(zuco2path, task_name, sbj_name, tn),
                error_callback=func.error_back
            )
            working_thread.append(w)
        for _ in range(workernum):
            working_thread[_].get()


def data_prepare_zuco2(zuco2path="./rawdata/zuco2", workernum=4):
    eeg, text = [], []
    for task_name in os.listdir(zuco2path):
        if "NR" in task_name:
            tn = "NR"
        elif "TSR" in task_name:
            tn = "TSR"
        else:
            continue
        sbj_names = list(os.listdir(f"{zuco2path}/{task_name}/Raw data"))
        for i in range(0, len(sbj_names), workernum):
            sbj_names_ = sbj_names[i:i + workernum]

            func = Func4MultiProcessing()
            mutipro_start(func, zuco2path, task_name, sbj_names_, tn)

    for tmpfile in os.listdir("./tmp"):
        if "zuco2" not in tmpfile:
            continue
        dic = pickle.load(open(f"./tmp/{tmpfile}", "rb"))
        eeg_seg, text_ = dic["eeg_seg"], dic["text_"]
        eeg += eeg_seg
        text += text_

    dic = {"eeg": eeg, "text": text}
    pickle.dump(dic, open("./data/zuco2.pkl", "wb"))
    print("---==Save [./data/zuco2.pkl]==---")